#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Description: 数据增强，取patch
@Date:2021/12/17 15:50:18
@Author: ljt
'''

import os
from numpy.lib.npyio import save
import paddle
import numpy as np
import glob 
from patchify import patchify, unpatchify
from matplotlib import image as mpimg
from tqdm import tqdm


dataset_root = "home/aistudio/data/data121840/dehw_train_dataset"

out_path = "/home/aistudio/data/data121840/dehw_train_dataset_patches"
os.makedirs(out_path, exist_ok=True)
os.makedirs(out_path + "/images", exist_ok=True)
os.makedirs(out_path + "/gts", exist_ok=True)

# get the img path list
im1_list = glob.glob(os.path.join(dataset_root, "images", "*.jpg"))
im2_list = glob.glob(os.path.join(dataset_root, "gts", "*.jpg"))
im1_list.sort()
im2_list.sort()
assert len(im1_list) == len(im2_list)


print(len(im2_list))

def save_patch(im_patch, save_path, num):
    for i in range(im_patch.shape[0]):
        for j in range(im_patch.shape[1]):
            # print(im_patch[i][j][0].shape)
            new_save_path = save_path + '/img{}_{}_{}.jpg'.format(num, i, j)
            mpimg.imsave(new_save_path, im_patch[i][j][0])


for id in tqdm(range(len(im1_list))):
    print(id)
    im1 = mpimg.imread(im1_list[id])
    im2 = mpimg.imread(im2_list[id])
    print(im1.shape)
    assert im1.shape == im2.shape
    name1 = os.path.basename(im1_list[id])
    name2 = os.path.basename(im2_list[id])
    # print(name1, name2)
    # assert name1 == name2
    
    ## when use a new dataset, we shoule run the block behind.
    # if im1.shape[0] < 256 or im1.shape[1] < 256:
    #     print("The dataset has image smaller than size 256, you should add th
    # e linear interpolation function.")
    #     break
    
    im1_patches = patchify(image=im1, patch_size=(256, 256, 3), step=128)
    im2_patches = patchify(image=im2, patch_size=(256, 256, 3), step=128)
    assert im1_patches.shape == im2_patches.shape
    save_patch(im1_patches, out_path + '/images', id)
    save_patch(im2_patches, out_path + '/gts', id)
    
